import cv2
import docx
import re
import os
from docx.table import Table
from docx.text.paragraph import Paragraph

from config_parser import ConfigParser
from utils import iter_block_items, get_pictures
from paddleocr import PPStructure, to_excel


def read_table(table):
    return [[cell.text for cell in row.cells] for row in table.rows]


class WordParser:
    def __init__(self, config, word_path):
        self.doc = docx.Document(word_path)
        self.cfg_txt = config.txt
        self.cfg_tbl = config.tbl
        self.cfg_pic = config.pic
        self.res = []
        self.parse()
        self.print_all()

    def parse(self):
        self.parse_pic()

        if not self.cfg_txt:
            return

        cfg = self.cfg_txt.pop(0)
        for block in iter_block_items(self.doc):
            if not cfg:
                return

            if isinstance(block, Paragraph):
                # print("text", [block.text])
                if self.parse_txt(block.text, cfg):
                    cfg = self.cfg_txt.pop(0) if self.cfg_txt else None
            elif isinstance(block, Table):
                # print("table", read_table(block))
                if self.parse_tbl(block, cfg):
                    cfg = self.cfg_txt.pop(0) if self.cfg_txt else None

    def parse_txt(self, text, cfg):
        if cfg[1] != 'text' or text.strip() == '':
            return False

        regex = re.sub('\[\w+\]', '(.*)', cfg[3])
        match = re.search(regex, text)
        if match:
            self.res.append((cfg[0], cfg[2], match.group(1)))
            return True
        else:
            return False

    def parse_tbl(self, block, cfg):
        if cfg[1] != 'table':
            return False

        table_info = [[cell.text for cell in row.cells] for row in block.rows]
        cfg_info = cfg[3]

        for (cfgi, tbli) in zip(cfg_info, table_info):
            if '[num]' in cfgi or '[str]' in cfgi:
                cfgi = ['[num]' if t == '[str]' else t for t in cfgi]
                self.res.append((cfg[0], cfgi[0], tbli[cfgi.index('[num]')]))
        return True

    def parse_pic(self):
        get_pictures(self.doc, r'./tmp')
        img_lst = os.listdir(r'./tmp')
        for cfg in self.cfg_pic:
            img_no = int(cfg[3].split('.')[0][-1]) - 1
            img_name = img_lst[img_no]
            exl_name = F'{cfg[2]}.xlsx'
            table_engine = PPStructure(show_log=False)
            img_path = os.path.join(r'./tmp', img_name)
            img = cv2.imread(img_path)
            result = table_engine(img)
            for res in result:
                if res['type'] == 'table':
                    to_excel(res['res']['html'], r'./data/doc/result/' + exl_name)

    def print_all(self):
        with open(r'./data/doc/result/word_info.txt', 'w') as f:
            para_old = ''
            for para, type, info in self.res:
                if para != para_old:
                    if para != '':
                        f.write('\n')
                    f.write(F'{para}\n{"-" * 20}\n')
                    para_old = para
                f.write(F'{type}: {info}\n')


if __name__ == '__main__':
    cfg_path = r'./data/doc/Test_doc_temp.docx'
    word_path = r'./data/doc/Test_doc.docx'
    cfg = ConfigParser(cfg_path)
    parser = WordParser(cfg, word_path)
    print('Done.')
